import paddle
import paddle.nn as nn
import paddle.nn.functional as F


class CrossEntropyCriterion(nn.Layer):
    def __init__(self):
        super(CrossEntropyCriterion, self).__init__()

    def forward(self, predict, label, seq_len, mask):
        # 生成mask
        seq_len = seq_len.squeeze(1)
        cost = F.softmax_with_cross_entropy(logits=predict, label=label.unsqueeze(2), soft_label=False)
        masked_cost = cost.squeeze(2) * mask
        sum_tp = paddle.sum(masked_cost, axis=1)  # 沿着batch_size使用,每一句话的loss
        seq_cost = paddle.mean(sum_tp)  # 每一句话的平均loss
        return seq_cost
